"""
simple.py
"""


import random
import pygraphviz


import unittest


class SanityGraphTester(unittest.TestCase):
    

    def test_single_add(self):
        g = SimpleGraph()
        g.add_node(1)

        resul = [1]
        self.assertEqual(resul, g.nodes())


    def test_single_del(self):
        g = SimpleGraph()
        g.add_node(1)
        g.del_node(1)

        resul = []
        self.assertEqual(resul, g.nodes())


    def test_multiple_add(self):
        g = SimpleGraph()
        resul = [1,2,3]
        g.add_nodes(resul)

        self.assertEqual(resul, g.nodes())


    def test_multiple_del(self):
        g = SimpleGraph()
        resul = [1,2,3]
        g.add_nodes(resul)
        g.del_nodes([1,3])

        self.assertEqual([2], g.nodes())


    def test_add_edge(self):
        g = SimpleGraph()
        g.add_nodes([1,2])
        g.add_edge(1,2)

        resul = [(1,2)]
        self.assertEqual(resul, g.edges())


    def test_del_edge(self):
        g = SimpleGraph()
        g.add_nodes([1,2,3,4])
        g.add_edge(1,2)
        g.add_edge(1,4)
        g.del_edge(1,2)

        resul = [(1,4)]
        self.assertEqual(resul, g.edges())


    def test_del_edge_reverse_order(self):
        g = SimpleGraph()
        g.add_nodes([1,2,3,4])
        g.add_edge(1,2)
        g.add_edge(1,4)
        g.del_edge(2,1)

        resul = [(1,4)]
        self.assertEqual(resul, g.edges())


    def test_nodes_after_del_with_edge(self):
        g = SimpleGraph()
        g.add_nodes([1,2,3,4])
        g.add_edge(1,2)
        g.add_edge(3,4)
        g.del_node(1)

        self.assertEqual([2,3,4], g.nodes())


    def test_edges_after_del_with_edge(self):
        g = SimpleGraph()
        g.add_nodes([1,2,3,4])
        g.add_edge(1,2)
        g.add_edge(3,4)
        g.del_node(1)

        self.assertEqual([(3,4)], g.edges())
        



class GraphError(Exception):
    pass


class NodeNotExistsError(GraphError):
    pass


class EdgeNotExistsError(GraphError):
    pass


class FileNotExistsError(GraphError):
    pass


class SimpleGraph():
    """Class with basic functionality for managing undirected graphs."""


    def __init__(self, name="simple-graph"):
        """Initialize a graph."""
        self.name = name
        self.adyacency  = {}
        self.conections = {}


    def add_node(self, node):
        """Add given node to the graph."""
        if not self.has_node(node):
            self.adyacency[node] = []


    def add_nodes(self, nodes):
        """Add given nodes to the graph."""
        for node in nodes:
            self.add_node(node)


    def del_node(self, node):
        """Remove a node from the graph."""
        if not self.has_node(node):
            raise NodeNotExistsError

        neighbors = list(self.neighbors(node))

        for neighbor in neighbors:
            self.del_edge(neighbor, node)
        
        del(self.adyacency[node])


    def del_nodes(self, nodes):
        """Remove nodes from the graph."""
        for node in nodes:
            self.del_node(node)


    def add_edge(self, source, target, weight=0):
        """Add an edge to the graph connecting nodes source and
        target."""
        if not self.has_node(source) or not self.has_node(target):
            raise NodeNotExistsError()

        if not self.has_edge(source, target):
            self.adyacency[source].append(target)
            self.adyacency[target].append(source)

            self.conections[(source, target)] = weight

    
    def del_edge(self, source, target):
        """Remove the edge (source, target) from the graph."""
        if not self.has_edge(source, target):
            raise EdgeNotExistsError()

        self.adyacency[source].remove(target)
        self.adyacency[target].remove(source)
        
        try:
            del(self.conections[(source, target)])
        except KeyError:
            del(self.conections[(target, source)])


    def nodes(self):
        """Return all the nodes in the graph."""
        return self.adyacency.keys()
        pass


    def edges(self):
        """Return all edges in the graph."""
        return self.conections.keys()


    def weight(self, edge):
        """Return the weight of an edge."""
        if not self.has_edge(edge[0], edge[1]):
            raise EdgeNotExistsError()

        try:
            return self.conections[edge]
        except KeyError:
            new_edge = (edge[1], edge[0])
            return self.conections[new_edge]


    def neighbors(self, node):
        """Return all nodes that are directly accessible from given
        node."""
        if not self.has_node(node):
            raise NodeNotExistsError()
        
        return self.adyacency[node]


    def has_node(self, node):
        """Return whether the requested node exists."""
        return node in self.nodes()

    
    def has_edge(self, source, target):
        """Return whether the requested edge exists."""
        p = (source, target)
        q = (target, source)

        return p in self.edges() or q in self.edges()
               

    def order(self, node):
        """Return the order of the given node."""
        if not self.has_node(node):
            raise NodeNotExistsError()

        return len(self.neighbors(node))


    def read(self, file):
        """Read a graph from a file in dot format."""
        try:
            graph = pygraphviz.AGraph(file)
        except IOError:
            raise FileNotExistsError()
        
        self.adyacency.clear()
        self.conections.clear()

        self.add_nodes(graph.nodes())

        for edge in graph.edges():
            weight = int(edge.attr["label"])
            self.add_edge(edge[0], edge[1], weight)


    def create_dot(self):
        """Create a dot format representation of the graph."""
        graph = pygraphviz.AGraph(name=self.name)

        graph.add_nodes_from(self.nodes())

        for edge in self.edges():
            weight = str(self.weight(edge))
            graph.add_edge(edge[0], edge[1], label=weight)

        return graph


    def write(self, name=None):
        """Write the graph to a file in dot format."""
        if not name:
            name  = self.name + ".dot"

        graph = self.create_dot()
        graph.write(name)
    

    def draw(self, nodes=None, name=None, color="red"):
        """Create an image representation of the graph."""
        if not name:
            name  = self.name
            
        name += ".png"

        graph = self.create_dot()

        for node in graph.nodes():
            location = node[1] + "," + node[4]
            node.attr["pos"] = location

        graph.graph_attr.update(overlap = "false")
        graph.node_attr.update(shape="circle")
        graph.node_attr.update(fontname = "Courier New")
        graph.edge_attr.update(fontname = "Courier New")

        if nodes:
            self.colorize(graph, nodes, color)
        
        graph.draw(path=name, format="png", prog="neato")


    def colorize(self, graph, nodes, color):
        """Paint the selected nodes in the graph."""
        for node in nodes:
            for guess in graph.nodes():
                if str(node) == guess:
                    guess.attr["color"] = color
                    guess.attr["style"] = "filled"


    def __str__(self):
        """Return the dot format representation of the graph."""
        graph = self.create_dot()
        return graph.string()


    def __len__(self):
        """Return the order of the graph."""
        return len(self.adyacency)



if __name__ == '__main__':
    unittest.main()
